import cv2
import numpy as np
import os
from os.path import exists
from imutils import paths
import pickle
from tqdm import tqdm
from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input
from tensorflow.keras.preprocessing import image
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 定义函数get_size,获取文件大小,返回文件大小的光字节的浮点数表示
def get_size(file):
    """ 
    获取指定文件的大小（以MB为单位）

    参数：
    file (str): 文件的路径

    返回：
    float: 文件大小（MB）
    """
    # TODO 填上函数的内容
    if not os.path.exists(file):
        return 0.0
    size_in_bytes = os.path.getsize(file)
    size_in_mb = size_in_bytes / (1024 * 1024)
    return size_in_mb

# 定义函数createXY,用于从图像创建特征(X)和标签(y)
def createXY(train_folder, dest_folder, method='vgg', batch_size=64):
    """  
    从训练文件夹中读取图像，生成特征向量和标签，并将它们保存到目标文件夹中  

    参数：
    train_folder (str): 包含训练图像的文件夹路径  
    dest_folder (str): 用于保存X和y文件的目标文件夹路径  
    method (str): 特征提取方法，'vgg' 或 'flat'
    batch_size (int): 批处理大小

    返回：
    (X, y) (tuple): 包含特征向量和标签的元组  
    """
    x_file_path = os.path.join(dest_folder, "X.pkl")  # 设置X文件的路径,用于保存特征数据
    y_file_path = os.path.join(dest_folder, "y.pkl")  # 设置y文件的路径,用于保存标签数据

    # 如果 X 和 y 已经存在,则直接读取,不再重新构建
    if os.path.exists(x_file_path) and os.path.exists(y_file_path):
        logging.info("X和y已经存在,直接读取")  # 提示用户X和y文件已经存在
        logging.info(f"X文件大小:{get_size(x_file_path):.2f}MB")  # 打印X文件的大小
        logging.info(f"y文件大小:{get_size(y_file_path):.2f}MB")  # 打印y文件的大小

        # TODO: 用 pickle 读取X和y文件
        with open(x_file_path, 'rb') as f:
            X = pickle.load(f)
        with open(y_file_path, 'rb') as f:
            y = pickle.load(f)
        
        return X, y

    logging.info("读取所有图像,生成X和y")  # 提示用户开始读取图像并生成X和y
    image_paths = list(paths.list_images(train_folder))  # 获取训练文件夹中所有图像的路径

    X = []  # 初始化X列表,用于存储特征
    y = []  # 初始化y列表,用于存储标签

    # 根据传入的方法选择不同的模型
    if method == 'vgg':
        model = VGG16(weights='imagenet', include_top=False, pooling="max")  # 加载VGG16模型,不包括顶层,使用最大池化
        logging.info("完成构建 VGG16 模型")  # 提示用户VGG16模型构建完成
    elif method == 'flat':
        model = None  # 如果方法为'flat',不使用任何预训练模型

    # 计算需要的批次数
    num_batches = len(image_paths) // batch_size + (1 if len(image_paths) % batch_size else 0)

    # 使用进度条对批次进行循环处理
    for idx in tqdm(range(num_batches), desc="读取图像"):
        batch_images = []  # 初始化存储批次图像的列表
        batch_labels = []  # 初始化存储批次标签的列表

        start = idx * batch_size  # 计算批次开始的索引
        end = min((idx + 1) * batch_size, len(image_paths))  # 计算批次结束的索引

        # 对于每个批次中的图像
        for i in range(start, end):
            image_path = image_paths[i]  # 获取图像路径
            if method == 'vgg':
                img = image.load_img(image_path, target_size=(224, 224))  # 加载图像并调整大小到224x224
                img = image.img_to_array(img)  # 将图像转换为数组
            elif method == 'flat':
                # TODO: 以灰度模式读取图像
                img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
                # TODO: 调整图像大小到32x32
                img = cv2.resize(img, (32, 32))
            batch_images.append(img)  # 将图像数据添加到批次图像列表

            # 从文件中获取标签
            # TODO: 从文件名中解析出字符串标签
            filename = os.path.basename(image_path)
            # TODO: 如果标签是'dog'则为1,否则为0
            label = 1 if 'dog' in filename.lower() else 0
            batch_labels.extend([label])  # 将标签添加到批次标签列表

        batch_images = np.array(batch_images)  # 将批次图像列表转换为numpy数组
        if method == 'vgg':
            batch_images = preprocess_input(batch_images)  # 对批次图像进行预处理
            batch_pixels = model.predict(batch_images, verbose=0)  # 使用VGG16模型进行预测
        else:
            # TODO: 将 batch_images 展平，注意，要展平的不只是一张图像
            batch_pixels = batch_images.reshape((batch_images.shape[0], -1))

        X.extend(batch_pixels)  # 将处理后的图像特征添加到X列表
        y.extend(batch_labels)  # 将标签添加到y列表

    logging.info(f"X.shape: {np.shape(X)}")  # 打印X的形状
    logging.info(f"y.shape: {np.shape(y)}")  # 打印y的形状

    # 将X和y分别序列化到文件
    # TODO: 用 pickle 保存X和y文件
    with open(x_file_path, 'wb') as f:
        pickle.dump(X, f)
    with open(y_file_path, 'wb') as f:
        pickle.dump(y, f)

    return X, y  # 返回构建的X和y